from langchain_core.messages import AIMessage
from langchain_core.messages import AnyMessage, HumanMessage
from langgraph.constants import START
from langgraph.graph import StateGraph
from typing_extensions import TypedDict, Annotated


def add(left, right):
    """Can also import `add` from the `operator` built-in."""
    return left + right


class State(TypedDict):
    messages: Annotated[list[AnyMessage], add]
    extra_field: int


def node(state: State):
    new_message = AIMessage("Hello!")
    return {"messages": [new_message], "extra_field": 10}


if __name__ == '__main__':
    graph = StateGraph(State).add_node(node).add_edge(START, "node").compile()
    result = graph.invoke({"messages": [HumanMessage("Hi")]})
    for message in result["messages"]:
        message.pretty_print()
